﻿import numpy as np
import colour
from colour.characterisation import CCS_COLOURCHECKERS
from colour.quality import colour_fidelity_index_ANSIIESTM3015

def calculate_tm30_metrics(test_spd, reference_spd = 'D65') :
    """
    计算光源的IES TM - 30 - 15指标(Rf和Rg)

    参数 :
    test_spd -- 待测光源的光谱功率分布(SpectralDistribution对象)
    reference_spd -- 参考光源(默认为D65), 可以是字符串或SpectralDistribution对象

    返回 :
Rf, Rg, 以及完整的TM - 30报告
"""
# 获取99个测试色样的光谱反射率数据
colour_checker = CCS_COLOURCHECKERS['99 Color Evaluation Samples']

# 计算TM - 30指标
report = colour_fidelity_index_ANSIIESTM3015(
    test_spd,
    additional_data = True,
    samples = colour_checker,
    reference_spd = reference_spd
)

# 提取关键指标
Rf = report.colour_fidelity_index
Rg = report.colour_gamut_index

return Rf, Rg, report

# ===================== 示例使用 =====================
if __name__ == "__main__":
# 1. 获取标准光源数据
from colour import SDS_ILLUMINANTS

# 2. 测试不同光源
test_lights = {
    'LED 3000K': SDS_ILLUMINANTS['LED-B1'],  # 暖白LED
    'Fluorescent F11': SDS_ILLUMINANTS['F11'],  # 荧光灯
    'Halogen': SDS_ILLUMINANTS['A'],  # 卤素灯
    'D65': SDS_ILLUMINANTS['D65']  # 参考光源
}

# 3. 计算并显示结果
print("{:<15} {:<10} {:<10} {:<10}".format("光源", "Rf", "Rg", "质量评价"))
print("=" * 45)

for name, spd in test_lights.items() :
    Rf, Rg, report = calculate_tm30_metrics(spd)

    # 质量评价
    quality = "优" if Rf >= 90 else "良" if Rf >= 80 else "中" if Rf >= 70 else "差"

    print("{:<15} {:<10.1f} {:<10.1f} {:<10}".format(name, Rf, Rg, quality))

    # 4. 可选：生成详细报告并可视化
    # 选择一个光源生成详细报告
    Rf, Rg, report = calculate_tm30_metrics(SDS_ILLUMINANTS['LED-B1'])

    # 打印详细报告
    print("\n详细报告:")
    print(report)

    # 可视化色域图
    report.plot_colour_vector_graphic()

    # 可视化色保真度图
    report.plot_colour_fidelity_graphic()
